import pefile
import os
import os
from typing import Dict, Optional
from utils.logger import setup_logger

logger = setup_logger()

class PEAnalyzer:
    """PE文件分析器"""
    
    def __init__(self):
        self.logger = setup_logger()
    
    def is_pe_file(self, file_path: str) -> bool:
        """
        检查文件是否为PE格式文件
        
        Args:
            file_path: 文件路径
            
        Returns:
            bool: 是否为PE文件
        """
        try:
            # 检查文件扩展名
            ext = os.path.splitext(file_path)[1].lower()
            if ext not in ['.exe', '.dll']:
                return False
            
            # 简单的文件头检查
            with open(file_path, 'rb') as f:
                header = f.read(2)
                if len(header) < 2 or header != b'MZ':
                    return False
            
            return True
        except Exception as e:
            self.logger.error(f"检查PE文件失败: {file_path}, 错误: {str(e)}")
            return False
    
    def extract_exe_info(self, file_path: str) -> Optional[Dict[str, str]]:
        """
        从可执行文件中提取软件详细信息
        
        Args:
            file_path: 可执行文件路径
            
        Returns:
            Dict: 提取的软件信息字典，失败返回None
        """
        try:
            if not self.is_pe_file(file_path):
                self.logger.warning(f"文件不是有效的PE文件: {file_path}")
                return None
            
            pe = pefile.PE(file_path)
            info = {}
            
            # 提取版本信息
            if hasattr(pe, 'FileInfo'):
                for fileinfo in pe.FileInfo:
                    if fileinfo.Key == b'StringFileInfo':
                        for st in fileinfo.StringTable:
                            # 将字节字符串转换为普通字符串
                            for key, value in st.entries.items():
                                try:
                                    str_key = key.decode('utf-8') if isinstance(key, bytes) else key
                                    str_value = value.decode('utf-8') if isinstance(value, bytes) else value
                                    info[str_key] = str_value
                                except UnicodeDecodeError:
                                    # 尝试使用其他编码
                                    try:
                                        str_key = key.decode('gbk') if isinstance(key, bytes) else key
                                        str_value = value.decode('gbk') if isinstance(value, bytes) else value
                                        info[str_key] = str_value
                                    except:
                                        continue
            
            # 清理和标准化信息
            return self._standardize_info(info)
            
        except Exception as e:
            self.logger.error(f"提取PE文件信息失败: {file_path}, 错误: {str(e)}")
            return None
    
    def _standardize_info(self, raw_info: Dict[str, str]) -> Dict[str, str]:
        """
        标准化提取的软件信息，统一键名
        
        Args:
            raw_info: 原始提取的信息
            
        Returns:
            Dict: 标准化后的信息
        """
        standardized = {}
        
        # 映射表，处理不同的键名变体
        key_mapping = {
            'ProductName': ['ProductName', '产品名称'],
            'FileDescription': ['FileDescription', '文件说明'],
            'CompanyName': ['CompanyName', '公司名称', '发布者', 'Publisher'],
            'FileVersion': ['FileVersion', '文件版本', '版本'],
            'LegalCopyright': ['LegalCopyright', '版权信息'],
            'OriginalFilename': ['OriginalFilename', '原始文件名']
        }
        
        for standard_key, possible_keys in key_mapping.items():
            for possible_key in possible_keys:
                for raw_key, value in raw_info.items():
                    if possible_key.lower() == raw_key.lower() and value.strip():
                        standardized[standard_key] = value.strip()
                        break
        
        return standardized
    
    def get_software_info_from_exe(self, file_path: str) -> Dict[str, Optional[str]]:
        """
        从可执行文件中提取结构化的软件信息
        
        Args:
            file_path: 可执行文件路径
            
        Returns:
            Dict: 结构化的软件信息
        """
        exe_info = self.extract_exe_info(file_path)
        if not exe_info:
            return {
                'name': None,
                'version': None,
                'publisher': None,
                'description': None
            }
        
        # 构建返回的软件信息
        software_info = {
            'name': exe_info.get('ProductName') or exe_info.get('FileDescription') or os.path.basename(file_path),
            'version': exe_info.get('FileVersion', 'Unknown'),
            'publisher': exe_info.get('CompanyName', 'Unknown'),
            'description': exe_info.get('FileDescription', '')
        }
        
        return software_info

# 全局实例，方便使用
pe_analyzer = PEAnalyzer()